package main

import (
	"fmt"
	"sync"
	"time"
)

type Unit struct {
	FirstTimeAt time.Time
	Count       int
	Next        *Unit
}

func NewUnits(c int) []*Unit {
	t := make([]*Unit, c)
	if c <= 1 {
		return t
	}
	for i := 0; i < c; i++ {
		t[i] = new(Unit)
	}
	for i := 0; i < c-1; i++ {
		t[i].Next = t[i+1]
	}
	t[c-1].Next = t[0]
	return t
}

type Limiter struct {
	//窗口内最多请求数量
	MaxInflight int
	//窗口间隔
	Duration time.Duration
	//窗口数量
	Split       int
	Units       []*Unit
	CurrentUnit *Unit
	lock        sync.Mutex
}

func NewLimiter(max int) Limiter {
	//5秒 ，每1秒一个单元
	return Limiter{
		MaxInflight: max,
		//滑动窗口间隔
		Duration: time.Second,
		Split:    5,
		Units:    NewUnits(5),
		lock:     sync.Mutex{},
	}
}

func (l *Limiter) Allow() bool {
	l.lock.Lock()
	defer l.lock.Unlock()
	if l.CurrentUnit == nil {
		l.CurrentUnit = l.Units[0]
	}
	if time.Now().Sub(l.CurrentUnit.FirstTimeAt) > time.Second {
		l.CurrentUnit = l.CurrentUnit.Next
		l.CurrentUnit.FirstTimeAt = time.Now()
		l.CurrentUnit.Count = 0
	}
	l.CurrentUnit.Count++
	inflight := 0
	for _, u := range l.Units {
		inflight += u.Count
	}
	if inflight >= l.MaxInflight {
		return false
	}
	return true

}

func main() {
	limiter := NewLimiter(10)
	start := time.Now()
	//some func or operation
	for i := 0; i <= 100; i++ {
		//限制10个请求，每个请求耗时200ms，则在第2秒开始限制
		time.Sleep(time.Millisecond * 200)
		if !limiter.Allow() {
			cost := time.Since(start)
			fmt.Printf("cost=[%s]", cost)
			fmt.Println(limiter.CurrentUnit.FirstTimeAt)
			break
		}
	}
}
